# -*- coding: utf-8 -*-
import pandas as pd
import matplotlib.pyplot as plt

import plotly.offline as py
import plotly.tools as tls
import plotly.figure_factory as FF
import plotly.graph_objs as go

class plotBase(object):
    def __init__(self):
        pass

    def _to_df(self, data):
        df = pd.DataFrame(data)
        if df.empty:
            return df
        df.set_index('datetime', inplace=True)
        df.index = pd.DatetimeIndex(df.index)
        return df


# 画图类, 这里用的是plotly. 也可以用其它的工具, 直接继承plotBase写新类就可以了
class plotly(plotBase):
    def __init__(self, context):
        # 画图首先需要拿到所有的基础的数据, 然后转换成dataframe格式
        self.context = context
        self.holdings = context.portfolio.all_holdings
        self.holdings_df = self._to_df(self.holdings)
        self.positions = context.portfolio.all_positions
        self.positions_df = self._to_df(self.positions)
        self.stocks_data = context.data_handler.latest_symbol_data
        self.stocks = self.context.stocks
        self.stocks_data_df = {}
        for i in self.stocks:
            self.stocks_data_df[i] = self._to_df(self.stocks_data[i])

        self.data = []

    def plot(self):
        # 根据画图的api来画的, 具体可以看官网的文档.
        for i in self.stocks_data_df:
            if self.stocks_data_df[i].empty:
                continue
            # print(self.stocks_data_df[i])
            p_stock = go.Scatter(x=self.stocks_data_df[i].index, y=self.stocks_data_df[i]['close'], xaxis='x', yaxis='y2', name=i)
            self.data.append(p_stock)

        for i in self.positions_df:
            p_position = go.Scatter(x=self.positions_df.index, y=self.positions_df[i], xaxis='x', yaxis='y', name=i + "_position")
            self.data.append(p_position)

        p_commission = go.Scatter(x=self.holdings_df.index, y=self.holdings_df['commission'], xaxis='x', yaxis='y3', name='commission')
        self.data.append(p_commission)

        p_cash = go.Scatter(x=self.holdings_df.index, y=self.holdings_df['cash'], xaxis='x', yaxis='y4', name='cash')
        self.data.append(p_cash)

        p_total = go.Scatter(x=self.holdings_df.index, y=self.holdings_df['total'], xaxis='x', yaxis='y4', name='total')
        self.data.append(p_total)

        layout = go.Layout(
            xaxis = dict(
                domain = [0, 1]
            ),
            yaxis = dict(
                domain = [0, 0.15],
                scaleanchor = 'x',
                title = 'positions'
            ),
            yaxis2 = dict(
                domain = [0.15, 0.65],
                scaleanchor='x',
                title = 'stocks'
            ),
            yaxis3 = dict(
                domain = [0.65, 0.85],
                scaleanchor='x',
                title = 'commissions'
            ),
            yaxis4 = dict(
                domain = [0.85, 1],
                scaleanchor='x',
                title = 'money'
            )

        )

        fig = go.Figure(data=self.data, layout=layout)
        py.plot(fig, filename='ploty.html', validate=False)